#!/usr/bin/env Rscript

# 03_annotation_differential_expression.R
# 细胞注释和差异表达分析
# 输入: seurat_obj_clustered.rds
# 输出: 注释后的对象、差异表达基因、标记基因可视化

library(Seurat)
library(ggplot2)
library(dplyr)
library(patchwork)
library(clustree)

# 设置随机种子
set.seed(123)

cat("=== 细胞注释和差异表达分析开始 ===\n")

# 读取聚类后的数据
cat("读取聚类后的Seurat对象...\n")
seurat_obj <- readRDS("output/seurat_obj_clustered.rds")

# 寻找每个聚类的标记基因
cat("寻找每个聚类的标记基因...\n")
cluster_markers <- FindAllMarkers(seurat_obj, 
                                 only.pos = TRUE, 
                                 min.pct = 0.25, 
                                 logfc.threshold = 0.25)

# 保存所有标记基因
write.csv(cluster_markers, "output/all_cluster_markers.csv", row.names = FALSE)

# 提取每个聚类的前5个标记基因
top_markers <- cluster_markers %>%
  group_by(cluster) %>%
  slice_max(n = 5, order_by = avg_log2FC)

write.csv(top_markers, "output/top_cluster_markers.csv", row.names = FALSE)

# 可视化标记基因
cat("生成标记基因可视化...\n")

# 热图展示前5个标记基因
top5_heatmap <- DoHeatmap(seurat_obj, features = top_markers$gene) + 
  theme(axis.text.y = element_text(size = 8))
ggsave("output/top_markers_heatmap.png", top5_heatmap, 
       width = 12, height = 10, dpi = 300)

# 点图展示前3个标记基因
top3_markers <- cluster_markers %>%
  group_by(cluster) %>%
  slice_max(n = 3, order_by = avg_log2FC)

dot_plot <- DotPlot(seurat_obj, features = unique(top3_markers$gene)) + 
  RotatedAxis() +
  theme(axis.text.x = element_text(angle = 45, hjust = 1))
ggsave("output/top_markers_dotplot.png", dot_plot, 
       width = 14, height = 8, dpi = 300)

# 手动细胞类型注释（基于已知标记基因）
cat("执行细胞类型注释...\n")

# 常见的PBMC标记基因
feature_plots <- FeaturePlot(seurat_obj, 
                            features = c("CD3D", "CD3E", "CD4", "CD8A", 
                                        "CD14", "FCGR3A", "MS4A1", 
                                        "NKG7", "GNLY", "PPBP"),
                            reduction = "umap", 
                            ncol = 3)

ggsave("output/cell_markers_umap.png", feature_plots, 
       width = 15, height = 12, dpi = 300)

# 基于标记基因表达进行细胞类型注释
new_cluster_ids <- c(
  "0" = "CD4+ T cells",
  "1" = "CD14+ Monocytes", 
  "2" = "CD8+ T cells",
  "3" = "B cells",
  "4" = "FCGR3A+ Monocytes",
  "5" = "NK cells",
  "6" = "Dendritic cells",
  "7" = "Megakaryocytes"
)

# 应用细胞类型注释
names(new_cluster_ids) <- levels(seurat_obj)
seurat_obj <- RenameIdents(seurat_obj, new_cluster_ids)

# 将细胞类型信息添加到metadata
seurat_obj$cell_type <- Idents(seurat_obj)

# 保存注释后的对象
saveRDS(seurat_obj, "output/seurat_obj_annotated.rds")

# 可视化注释结果
cat("生成注释结果可视化...\n")

# 注释后的UMAP图
annotated_umap <- DimPlot(seurat_obj, 
                         reduction = "umap", 
                         label = TRUE, 
                         pt.size = 0.5) + 
  NoLegend() +
  ggtitle("Cell Type Annotation")

ggsave("output/annotated_umap.png", annotated_umap, 
       width = 10, height = 8, dpi = 300)

# 带图例的版本
annotated_umap_legend <- DimPlot(seurat_obj, 
                                reduction = "umap", 
                                label = FALSE, 
                                pt.size = 0.5) +
  ggtitle("Cell Type Annotation")

ggsave("output/annotated_umap_legend.png", annotated_umap_legend, 
       width = 10, height = 8, dpi = 300)

# 细胞类型比例图
cell_type_counts <- table(seurat_obj$cell_type)
cell_type_df <- data.frame(
  CellType = names(cell_type_counts),
  Count = as.numeric(cell_type_counts)
)

cell_type_plot <- ggplot(cell_type_df, aes(x = reorder(CellType, -Count), y = Count, fill = CellType)) +
  geom_bar(stat = "identity") +
  theme_minimal() +
  theme(axis.text.x = element_text(angle = 45, hjust = 1)) +
  labs(x = "Cell Type", y = "Number of Cells", title = "Cell Type Distribution") +
  scale_fill_brewer(palette = "Set3")

ggsave("output/cell_type_distribution.png", cell_type_plot, 
       width = 10, height = 6, dpi = 300)

# 保存细胞类型统计
write.csv(cell_type_df, "output/cell_type_statistics.csv", row.names = FALSE)

# 差异表达分析示例：比较特定细胞类型
cat("执行特定细胞类型的差异表达分析...\n")

# 示例：比较T细胞和单核细胞
if ("CD4+ T cells" %in% unique(seurat_obj$cell_type) & 
    "CD14+ Monocytes" %in% unique(seurat_obj$cell_type)) {
  
  t_vs_monocyte_markers <- FindMarkers(seurat_obj,
                                      ident.1 = "CD4+ T cells",
                                      ident.2 = "CD14+ Monocytes",
                                      min.pct = 0.25)
  
  write.csv(t_vs_monocyte_markers, "output/t_cell_vs_monocyte_markers.csv")
  
  # 可视化前几个差异基因
  top_de_genes <- rownames(t_vs_monocyte_markers)[1:6]
  de_feature_plot <- FeaturePlot(seurat_obj, 
                                features = top_de_genes,
                                reduction = "umap", 
                                ncol = 3)
  
  ggsave("output/de_genes_umap.png", de_feature_plot, 
         width = 12, height = 8, dpi = 300)
}

# 生成分析报告
cat("生成分析摘要...\n")
analysis_summary <- data.frame(
  Metric = c("Total Cells", "Number of Clusters", "Number of Cell Types"),
  Value = c(ncol(seurat_obj), 
           length(unique(seurat_obj@meta.data$seurat_clusters)),
           length(unique(seurat_obj$cell_type)))
)

write.csv(analysis_summary, "output/analysis_summary.csv", row.names = FALSE)

cat("=== 细胞注释和差异表达分析完成 ===\n")
cat("输出文件:\n")
cat("- output/seurat_obj_annotated.rds: 注释后的Seurat对象\n")
cat("- output/all_cluster_markers.csv: 所有聚类标记基因\n")
cat("- output/top_cluster_markers.csv: 前5个标记基因\n")
cat("- output/annotated_umap*.png: 注释结果可视化\n")
cat("- output/cell_type_statistics.csv: 细胞类型统计\n")
cat("- output/cell_markers_umap.png: 标记基因表达图\n")